import torch
from torch import nn
from torchvision.datasets import ImageFolder  # 直接加载数据目录，进行分类
from torchvision.transforms import transforms  # 将图像进行一定处理的工具
from torch.utils.data import DataLoader

# ==================== 图像数据集处理 ===================
# 对单张图片进行处理
transform = transforms.Compose([
    transforms.ToTensor()  # 将加载的图像数据转化为pytorch.tensor类型
])

# ImageFoler将指定目录下的文件数据进行自动分类
dataset = ImageFolder("./data/numbers", transform=transform)

# 将上述数据集进行 将数据分批次、将数据进行随机打散
dataloader = DataLoader(dataset, batch_size=1000, shuffle=True)

# ==================== 搭建训练模型（卷积神经网络） ===================
# 训练手写数字数据集 （28，28，3） --> (6-8,6-8,n) 最终卷积之后的图像大小必须在 6-8 之间（重点）
model = nn.Sequential(
    nn.Conv2d(3, 16, 3, 1, 1),  # (28x28x16)
    nn.ReLU(),
    nn.MaxPool2d(2, 2),  # (14x14x16)
    nn.Conv2d(16, 32, 3, 1, 1),  # (14x14x32)
    nn.ReLU(),
    nn.MaxPool2d(2, 2),  # (7x7x32)
    nn.Flatten(),  # (1568,) 将上述的图像矩阵进行拉伸（展平）
    nn.Dropout(),  # 随机削减50% 的神经元数量（训练提速）
    nn.Linear(7 * 7 * 32, 1024),
    nn.ReLU(),
    nn.Linear(1024, 10),  # 10是分类的数量
    nn.LogSoftmax(dim=-1)  # 多分类激活函数
)
